#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Time          : 2024/4/14 18:12
# File          : snow_flake_id.py
# @Author       : MingTai
# @Department   : automation_qa

# https://github.com/Grokzen/redis-py-cluster

# 普通单服务redis
from redis.client import Redis
from redis.connection import ConnectionPool
from redis.sentinel import Sentinel
from rediscluster import RedisCluster

from app.utils.redis_client.redis_config import REDIS_HOST, REDIS_PASSWORD


class RedisClient:

    def __init__(self, host, port, password, db=0, max_connections=10, **connection_kwargs):
        if host and port and password:
            self.conn_pool = ConnectionPool(host=host, port=port, password=password, db=db,
                                            max_connections=max_connections, **connection_kwargs)
        else:
            print("redis's config: cluster_nodes or password not exist!")

    def get_redis_client(self):
        redis_pool = Redis(connection_pool=self.conn_pool)
        return redis_pool

    def get_key(self, key):
        return self.get_redis_client().get(key)

    def set_key(self, key, value, **kwargs):
        return self.get_redis_client().set(key, value, **kwargs)

    def delete_key(self, keys):
        return self.get_redis_client().delete(keys)

    def append(self, key, value):
        return self.get_redis_client().append(key, value)

    def lpush(self, key, value):
        self.get_redis_client().lpush(key, value)

    def rpush(self, key, value):
        self.get_redis_client().rpush(key, value)

    def lrange(self, key, start, end):
        return self.get_redis_client().lrange(key, start, end)


# redis 集群模式
class RedisClusterClient:

    def __init__(self, cluster_nodes, password, max_connections=200, **kwargs):
        if cluster_nodes and password:
            self.RedisClientPool = RedisCluster(startup_nodes=cluster_nodes, password=password, decode_responses=True,
                                                max_connections=max_connections, **kwargs)
        else:
            print("redis's config: cluster_nodes or password not exist!")

    def get_redis_client(self):
        return self.RedisClientPool

    def get_key(self, key):
        return self.RedisClientPool.get(key)

    def delete_key(self, keys):
        return self.RedisClientPool.delete(keys)


# redis 哨兵模式
class RedisSentinel:
    def __init__(self, sentinel_list, password, service_name, db=0):
        self.sentinel = Sentinel(sentinel_list, socket_timeout=0.5)
        self.service_name = service_name
        self.password = password
        self.db = db
        self.master = None
        self.slave = None

    def get_master_redis(self):
        return self.sentinel.discover_master(self.service_name)

    def get_slave_redis(self):
        return self.sentinel.discover_slaves(self.service_name)

    def get_master_client(self):
        if not self.master:
            self.master = self.sentinel.master_for(
                service_name=self.service_name,
                socket_timeout=0.5,
                password=self.password,
                db=self.db
            )
        return self.master

    def get_slave_client(self):
        if not self.slave:
            self.slave = self.sentinel.slave_for(
                service_name=self.service_name,
                socket_timeout=0.5,
                password=self.password,
                db=self.db
            )
        return self.slave

    def get_key(self, kek):
        self.get_slave_client().get(kek)

    def set_key(self, kek, value, **kwargs):
        self.get_slave_client().set(kek, value, **kwargs)

    def rpush(self, key, value):
        self.get_slave_client().rpush(key, value)

    def lpush(self, key, value):
        self.get_slave_client().lpush(key, value)

    def lrange(self, key, start, end):
        self.get_slave_client().lrange(key, start, end)


def create_redis_client_fun():
    return RedisClusterClient(REDIS_HOST, REDIS_PASSWORD).get_redis_client()


if __name__ == '__main__':
    # 集群模式
    # redisClust = RedisClusterClient(config_test.REDIS_HOST, REDIS_PASSWORD).get_redis_client()
    redisClust = RedisClusterClient(REDIS_HOST, REDIS_PASSWORD).get_redis_client()
    redisClust.set("foo", "bar")
    print(redisClust.get("foo"))
    # 哨兵模式
    # redisSent = RedisSentinel(sentinel_list=REDIS_HOST,
    #                           password=REDIS_PASSWORD,
    #                           service_name=REDIS_SENTINEL_NAME).get_master_client()
